package vnc

import (
	"encoding/xml"
	"fmt"
	"log"
	"net"
	"net/http"

	"github.com/gorilla/websocket"
	libvirt "libvirt.org/go/libvirt"
	"libvirt.org/go/libvirtxml"
)

// VNCProxy VNC代理服务器
type VNCProxy struct {
	upgrader websocket.Upgrader
}

// NewVNCProxy 创建新的VNC代理
func NewVNCProxy() *VNCProxy {
	return &VNCProxy{
		upgrader: websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool {
				// 允许所有来源的连接
				return true
			},
		},
	}
}

// HandleWebSocket 处理WebSocket连接
func (p *VNCProxy) HandleWebSocket(w http.ResponseWriter, r *http.Request, domainName string, vncHost string, vncPort int) {
	// 升级到WebSocket连接
	ws, err := p.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("Failed to upgrade to WebSocket: %v", err)
		return
	}
	defer ws.Close()

	log.Printf("WebSocket VNC proxy started for domain %s, connecting to %s:%d", domainName, vncHost, vncPort)

	// 连接到实际的VNC服务器
	vncAddr := fmt.Sprintf("%s:%d", vncHost, vncPort)
	vncConn, err := net.Dial("tcp", vncAddr)
	if err != nil {
		log.Printf("Failed to connect to VNC server %s: %v", vncAddr, err)
		ws.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf("Failed to connect to VNC server: %v", err)))
		return
	}
	defer vncConn.Close()

	// 发送连接成功的消息
	ws.WriteMessage(websocket.TextMessage, []byte("VNC WebSocket proxy connected"))

	// 启动双向数据传输
	errChan := make(chan error, 2)

	// 从WebSocket读取并写入到VNC
	go func() {
		for {
			messageType, message, err := ws.ReadMessage()
			if err != nil {
				errChan <- fmt.Errorf("WebSocket read error: %v", err)
				return
			}

			// 只转发二进制消息到VNC
			if messageType == websocket.BinaryMessage {
				_, err = vncConn.Write(message)
				if err != nil {
					errChan <- fmt.Errorf("VNC write error: %v", err)
					return
				}
			}
		}
	}()

	// 从VNC读取并写入到WebSocket
	go func() {
		buffer := make([]byte, 4096)
		for {
			n, err := vncConn.Read(buffer)
			if err != nil {
				errChan <- fmt.Errorf("VNC read error: %v", err)
				return
			}

			err = ws.WriteMessage(websocket.BinaryMessage, buffer[:n])
			if err != nil {
				errChan <- fmt.Errorf("WebSocket write error: %v", err)
				return
			}
		}
	}()

	// 等待错误发生
	err = <-errChan
	log.Printf("VNC proxy error: %v", err)
}

// GetVNCConnectionInfo 获取VNC连接信息
func GetVNCConnectionInfo(conn *libvirt.Connect, domainName string) (host string, port int, err error) {
	dom, err := conn.LookupDomainByName(domainName)
	if err != nil {
		return "", 0, fmt.Errorf("domain not found: %w", err)
	}
	defer dom.Free()

	// 获取域的XML描述
	xmlDesc, err := dom.GetXMLDesc(0)
	if err != nil {
		return "", 0, fmt.Errorf("failed to get domain XML: %w", err)
	}

	// 解析XML
	var domainDef libvirtxml.Domain
	if err := xml.Unmarshal([]byte(xmlDesc), &domainDef); err != nil {
		return "", 0, fmt.Errorf("failed to parse domain XML: %w", err)
	}

	// 查找VNC图形设备
	if domainDef.Devices != nil && len(domainDef.Devices.Graphics) > 0 {
		for _, graphic := range domainDef.Devices.Graphics {
			if graphic.VNC != nil {
				vnc := graphic.VNC
				host = "127.0.0.1" // 默认主机
				port = vnc.Port

				// 如果配置了监听地址，则使用配置的地址
				if vnc.Listen != "" {
					host = vnc.Listen
				}

				// 如果端口号为-1，则需要获取实际分配的端口号
				if vnc.Port == -1 {
					// 获取实际运行时信息
					display, err := getActualVNCDisplay(dom)
					if err != nil {
						log.Printf("Failed to get actual VNC display: %v", err)
					} else {
						port = 0 + display
						log.Printf("VNC display %d, port %d", display, port)
					}
				} else if vnc.Port >= 0 {
					// 标准VNC端口是5900 + 显示编号
					port = 0 + vnc.Port
				}

				return host, port, nil
			}
		}
	}

	return "", 0, fmt.Errorf("no VNC graphics device found")
}

// getActualVNCDisplay 获取运行时实际分配的VNC显示编号
func getActualVNCDisplay(dom *libvirt.Domain) (int, error) {
	// 获取域的XML描述（运行时状态）
	xmlDesc, err := dom.GetXMLDesc(libvirt.DOMAIN_XML_SECURE)
	if err != nil {
		return 0, err
	}

	// 解析XML
	var domainDef libvirtxml.Domain
	if err := xml.Unmarshal([]byte(xmlDesc), &domainDef); err != nil {
		return 0, err
	}

	// 查找VNC图形设备
	if domainDef.Devices != nil && len(domainDef.Devices.Graphics) > 0 {
		for _, graphic := range domainDef.Devices.Graphics {
			if graphic.VNC != nil && graphic.VNC.Port > 0 {
				// 返回实际的显示编号
				return graphic.VNC.Port, nil
			}
		}
	}

	return 0, fmt.Errorf("no VNC display found in runtime XML")
}
